#include "dllso/libwine.hpp"

#include "dllso/DllsoInit.hpp"
#include "misc/arch.hpp"
#include "misc/helper.hpp"
#include "rewine/LoggerSingleton.hpp"
#include "rewine/MapIndicator.hpp"

#include "wininc/pe.h"

#include <cassert>
#include <memory>

#define WINE_VER "WINE_1.0"
#define SYMVER(s) __asm__(".symver " #s ", " #s "@" WINE_VER)

SYMVER(__wine_main_argc);
SYMVER(__wine_main_argv);
SYMVER(__wine_main_wargv);
SYMVER(__wine_main_environ);
SYMVER(__wine_dll_register);

int __wine_main_argc = 0;
char **__wine_main_argv = nullptr;
wchar_t **__wine_main_wargv = nullptr;
char **__wine_main_environ = nullptr;

namespace rewine
{

static inline void *
rva2va(INT_PTR rva, void *base) {
    return (void *)(rva + (INT_PTR) base);
}

static inline INT_PTR
va2rva(void *va, void *base) {
    return (INT_PTR) va - (INT_PTR) base;
}

static inline INT_PTR
rebase_ptr(INT_PTR rva, long delta) {
    return rva + delta;
}

static inline DWORD
rebase_dword(DWORD rva, long delta) {
    return rva + delta;
}

template <typename T>
static inline T
align_to(T value, unsigned int mask) {
    return (value + mask) & ~(T)mask;
}

static inline void
fixup_rva_ptrs(void *array, void *base, unsigned int count) {
    // va2rva and INT_PTR (4/8) to DWORD (4)
    auto src = (void **)array;
    auto dst = (DWORD *)array;
    while (count--) {
        *dst = (*src) ? (INT_PTR)(*src) - (INT_PTR) base : 0;
        src++; dst++;
    }
}

static inline void
fixup_rva_dwords(DWORD *ptr, long delta, unsigned int count) {
    // rebase DWORD (4)
    // formula: delta = new - old; ret = (rva - old) + new;
    while (count--) {
        if (*ptr) *ptr += delta;
        ptr++;
    }
}

static inline void
fixup_rva_names(UINT_PTR *ptr, long delta) {
    // rebase POINTER (4/8)
    while (*ptr) {
        if (!IMAGE_SNAP_BY_ORDINAL(*ptr)) *ptr += delta;
        ptr++;
    }
}

static void
fixup_exports(IMAGE_EXPORT_DIRECTORY *dir, void *base, int delta) {
    fixup_rva_dwords(&dir->Name, delta, 1);
    fixup_rva_dwords(&dir->AddressOfFunctions, delta, 1);
    fixup_rva_dwords(&dir->AddressOfNames, delta, 1);
    fixup_rva_dwords(&dir->AddressOfNameOrdinals, delta, 1);
    fixup_rva_dwords((DWORD *) rva2va(dir->AddressOfNames, base), delta, dir->NumberOfNames);
    fixup_rva_ptrs(rva2va(dir->AddressOfFunctions, base), base, dir->NumberOfFunctions);
}

static void
fixup_imports(IMAGE_IMPORT_DESCRIPTOR *dir, void *base, int delta) {
    while (dir->Name) {
        fixup_rva_dwords(&dir->u.OriginalFirstThunk, delta, 1);
        fixup_rva_dwords(&dir->Name, delta, 1);
        fixup_rva_dwords(&dir->FirstThunk, delta, 1);
        if (dir->u.OriginalFirstThunk) fixup_rva_names((UINT_PTR *) rva2va(dir->u.OriginalFirstThunk, base), delta);
        if (dir->FirstThunk) fixup_rva_names((UINT_PTR *) rva2va(dir->FirstThunk, base), delta);
        dir++;
    }
}

static std::shared_ptr<MapIndicator>
fixup_dllso(const void *old_base) {
    int ret;

    const size_t page_mask = rewine::get_page_mask();

    auto ntOrigin = reinterpret_cast<const IMAGE_NT_HEADERS *>(old_base);
    auto szHeaders = ntOrigin->OptionalHeader.SizeOfHeaders;

    // module address must be aligned on 64K boundary
    auto new_base = (void *)(align_to((UINT_PTR) ntOrigin->OptionalHeader.ImageBase, 0xFFFFu));

    std::shared_ptr<MapIndicator> headersMap;
    ret = mmap_anonymous(new_base, szHeaders, false, headersMap);
    if (ret < 0) {
        return nullptr;
    }

    auto dos = (PIMAGE_DOS_HEADER) headersMap->next(sizeof(IMAGE_DOS_HEADER));
    if (!dos) {
        return nullptr;
    }
    dos->e_magic    = IMAGE_DOS_SIGNATURE;
    dos->e_cblp     = 0x90;
    dos->e_cp       = 3;
    dos->e_cparhdr  = (sizeof(*dos) + 0xf) / 0x10;
    dos->e_minalloc = 0;
    dos->e_maxalloc = 0xffff;
    dos->e_ss       = 0x0000;
    dos->e_sp       = 0x00b8;
    dos->e_lfarlc   = sizeof(*dos);
    dos->e_lfanew   = sizeof(*dos);

    auto nt = (PIMAGE_NT_HEADERS) headersMap->next(sizeof(IMAGE_NT_HEADERS));
    if (!nt) {
        return nullptr;
    }
    *nt = *ntOrigin;

    auto old_base_value = (UINT_PTR) old_base;
    auto new_base_value = (UINT_PTR) headersMap->front();
    auto delta = (long) -(new_base_value - old_base_value);  // why negative, by design?

    // SizeOfImage == _end - rva_base
    auto szImage = align_to(ntOrigin->OptionalHeader.SizeOfImage + delta, page_mask);

    nt->FileHeader.NumberOfSections                = 0;
    nt->OptionalHeader.BaseOfCode                  = 0;
#ifndef _WIN64
    nt->OptionalHeader.BaseOfData                  = 0;
#endif
    nt->OptionalHeader.SizeOfCode                  = 0;
    nt->OptionalHeader.SizeOfInitializedData       = 0;
    nt->OptionalHeader.SizeOfUninitializedData     = 0;
    nt->OptionalHeader.SizeOfImage                 = szImage;
    nt->OptionalHeader.ImageBase                   = new_base_value;

    fixup_rva_ptrs(&nt->OptionalHeader.AddressOfEntryPoint, new_base, 1);

    for (auto i = 0; i < nt->OptionalHeader.NumberOfRvaAndSizes; i++) {
        fixup_rva_dwords(&nt->OptionalHeader.DataDirectory[i].VirtualAddress, delta, 1);
    }

    PIMAGE_DATA_DIRECTORY dir;

    dir = &nt->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_EXPORT];
    if (dir->Size) {
        auto exports = (IMAGE_EXPORT_DIRECTORY *) rva2va(dir->VirtualAddress, new_base);
        fixup_exports(exports, new_base, delta);
    }

    dir = &nt->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT];
    if (dir->Size) {
        auto imports = (IMAGE_IMPORT_DESCRIPTOR *) rva2va(dir->VirtualAddress, new_base);
        fixup_imports(imports, new_base, delta);
    }

    return std::move(headersMap);
}

} // namespace rewine

void
__wine_dll_register(const void *nt, const char *filename) {
    rewine::g_Logger->debug(__FLF__, "dllso callback {}", filename);
    auto fixed = rewine::fixup_dllso(nt);
    rewine::dllso::putDllsoBase(filename, fixed);
}
